import pybullet as p
import time
import pybullet_data
import numpy as np
import matplotlib

matplotlib.use('Qt5Agg')
import matplotlib.pyplot as plt

"""
zc: z-axis intersection
Tsup: support time
sx, sy: walk parameters(vectors)
theta: changing degrees (not rad!) per step
xi, yi: initial CoM coordinate
px, py: initial CoM foot placement
vxi, vyi: initial CoM velocity
a, b: params of cost function
"""

n = 0  # step number
g = 9.81  # gravity: m/s^2
a = 10
b = 1
zc = 0.8
Tsup = 0.8
sx = [0.0, 0.3, 0.3, 0.3, 0.3, 0.0]
sy = [0.0032, 0.0032, 0.0032, 0.0032, 0.0032, 0.0032]
theta = [0, 0, 0, 0, 0, 0]

"""The step after the last step should be zero"""
# sx.append(0)
# sy.append(0)

"""Calculate basic paras"""
Tc = np.sqrt(zc / g)
C = np.cosh(Tsup / Tc)
S = np.sinh(Tsup / Tc)
D = a * (C - 1) ** 2 + b * (S / Tc) ** 2
radian = np.deg2rad(theta)
st = np.sin(radian)
ct = np.cos(radian)


def rotate(i):
    if i == 0:
        R = np.eye(2)
    else:
        R = np.array([[ct[i], -st[i]], [st[i], ct[i]]])
    return R


class Biped3D(object):
    def __init__(self):
        self.physicsClient = p.connect(p.GUI)  # or p.DIRECT for non-graphical version
        p.resetDebugVisualizerCamera(cameraDistance=0.5, cameraYaw=60,
                                     cameraPitch=-30, cameraTargetPosition=[0, 0, 0.2])
        p.setAdditionalSearchPath(pybullet_data.getDataPath())  # optionally

        self.ground = p.loadURDF("plane.urdf")
        self.robot = p.loadURDF("/yanshee/robots/yanshee.urdf", [0, 0, 0], [0, 0, 0, 1], useFixedBase=False)
        self.joints = self.get_joints()
        self.n_j = len(self.joints)
        self.simu_f = 1000  # Simulation frequency, Hz
        p.setTimeStep(1.0 / self.simu_f)
        self.motion_f = 2  # Controlled motion frequency, Hz
        self.q_vec = np.zeros(self.n_j)
        self.dq_vec = np.zeros(self.n_j)
        # self.q_mat = np.zeros((self.simu_f * 3, self.n_j))
        # self.q_d_mat = np.zeros((self.simu_f * 3, self.n_j))
        self.q_mat = np.zeros((self.simu_f * 9, self.n_j))  # 9*500 rows for simsyncs & 9 columns for all joints
        self.q_d_mat = np.zeros((self.simu_f * 9, self.n_j))
        # self.init_plot()

        """Foot placement paras"""
        self.px = 0
        self.py = 0
        self.xi = 0
        self.yi = 0
        self.vxi = 0
        self.vyi = 0
        self.px0 = self.px
        self.py0 = self.py
        """
        xt: x_tmp
        yt: y_tmp
        vxt: vx_tmp
        vyt: vy_tmp
        """
        self.xt = 0
        self.yt = 0
        self.vxt = 0
        self.vyt = 0

    def run(self):
        for i in range(int(5e3)):
            t = i / self.simu_f
            torque_array = self.controller(t)
            self.q_vec, self.dq_vec = self.step(torque_array)
            # if 0 == i % 20:
            #     self.update_plot()
            time.sleep(1 / self.simu_f)
        p.disconnect()

    def step(self, torque_array):
        self.set_motor_torque_array(torque_array)
        p.stepSimulation()
        self.q_mat[:-1] = self.q_mat[1:]
        self.q_mat[-1] = self.q_vec
        return self.get_joint_states()

    def get_joints(self):
        all_joints = []
        for j in range(p.getNumJoints(self.robot)):
            # Disable motor in order to use direct torque control.
            info = p.getJointInfo(self.robot, j)
            joint_type = info[2]
            if (joint_type == p.JOINT_PRISMATIC or joint_type == p.JOINT_REVOLUTE):
                all_joints.append(j)
                p.setJointMotorControl2(self.robot, j,
                                        controlMode=p.POSITION_CONTROL, force=0)
        joints = all_joints[0:]
        # print("Number of All Joints:", p.getNumJoints(self.robot))
        # print("Number of All Revolute Joints:", joints)
        return joints

    def get_joint_states(self):
        '''
        :return: q_vec: joint angle, dq_vec: joint angular velocity
        '''
        q_vec = np.zeros(self.n_j)
        dq_vec = np.zeros(self.n_j)
        for j in range(self.n_j):
            q_vec[j], dq_vec[j], _, _ = p.getJointState(self.robot, self.joints[j])
        return q_vec, dq_vec

    def set_motor_torque_array(self, torque_array=None):
        '''
        :param torque_array: the torque of [rHAA, rHFE, rKFE]
        :Specifically, H: Hip, K: Knee, AA:abduction/adduction, FE:flexion/extension
        '''
        if torque_array is None:
            torque_array = np.zeros(self.n_j)
        for j in range(len(self.joints)):
            p.setJointMotorControl2(self.robot, self.joints[j], p.TORQUE_CONTROL, force=torque_array[j])

    def controller(self, t, type='joint'):
        if 'joint' == type:
            return self.joint_controller(t)

    def joint_controller(self, t):

        # lf_pos = p.getLinkState(1, 11)[0]
        # rf_pos = p.getLinkState(1, 16)[0]
        # q_d_ctr = max(lf_pos[0], rf_pos[0])
        print("n")
        print(n)
        q_d_ctr, dq_d_ctr = self.cal_CoM_Trajectory(t)  # Note the CoM Trajectory
        # print("q_d_ctr")
        # print(q_d_ctr)
        # print("dq_d_ctr")
        # print(dq_d_ctr)
        amp = np.deg2rad(30)
        q_d_haa = 0
        q_d_hfe = amp * np.sin(2 * np.pi * self.motion_f * t)
        q_d_kfe = amp / 2.0 * (np.cos(2 * np.pi * self.motion_f * t) + 1)

        # q_d_vec = np.array([q_d_haa, q_d_hfe, q_d_kfe])
        q_d_vec = np.array(
            [q_d_ctr, q_d_haa, q_d_hfe, q_d_kfe, -q_d_hfe - q_d_kfe, q_d_haa, q_d_hfe, q_d_kfe, -q_d_hfe - q_d_kfe]
        )
        self.q_d_mat[:-1] = self.q_d_mat[1:]

        # print("q_d_mat[-1]")
        # print(self.q_d_mat[-1])
        # print("q_d_vec")
        # print(q_d_vec)
        self.q_d_mat[-1] = q_d_vec

        dq_d_haa = 0
        dq_d_hfe = 2 * np.pi * self.motion_f * amp * np.cos(2 * np.pi * self.motion_f * t)
        dq_d_kfe = -np.pi * self.motion_f * amp * np.sin(2 * np.pi * self.motion_f * t)
        # dq_d_vec = np.array([dq_d_haa, dq_d_hfe, dq_d_kfe])
        dq_d_vec = np.array(
            [dq_d_ctr, dq_d_haa, dq_d_hfe, dq_d_kfe, -dq_d_hfe - q_d_kfe, dq_d_haa, dq_d_hfe, dq_d_kfe,
             -dq_d_hfe - dq_d_kfe]
        )

        # PD controller gains
        # k_vec = np.array([100, 50, 50])
        # b_vec = np.array([0.05, 0.05, 0.05])
        k_vec = np.array([100, 100, 50, 50, 50, 100, 50, 50, 50])
        b_vec = np.array([0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05])
        # print(k_vec)
        # print(b_vec)
        # print(self.q_vec)
        # print(self.dq_vec)
        # print(q_d_vec)
        # print(dq_d_vec)
        return self.joint_PD_controller(k_vec, b_vec, self.q_vec, self.dq_vec, q_d_vec, dq_d_vec)

    def joint_PD_controller(self, k_vec, b_vec, q_vec, dq_vec, q_d_vec, dq_d_vec):
        return k_vec * (q_d_vec - q_vec) + b_vec * (dq_d_vec - dq_vec)

    # def init_plot(self):
    #     self.fig = plt.figure(figsize=(5, 9))
    #     joint_names = ['rHAA', 'rHFE', 'rKFE', ]
    #     self.q_d_lines = []
    #     self.q_lines = []
    #     for i in range(3):
    #         plt.subplot(3, 1, i + 1)
    #         q_d_line, = plt.plot(self.q_d_mat[:, i], '-')
    #         q_line, = plt.plot(self.q_mat[:, i], '--')
    #         self.q_d_lines.append(q_d_line)
    #         self.q_lines.append(q_line)
    #         plt.ylabel('q_{} (rad)'.format(joint_names[i]))
    #         plt.ylim([-1.5, 1.5])
    #     plt.xlabel('Simulation steps')
    #     self.fig.legend(['q_d', 'q'], loc='lower center', ncol=2, bbox_to_anchor=(0.49, 0.97), frameon=False)
    #     self.fig.tight_layout()
    #     plt.draw()

    # def update_plot(self):
    #     for i in range(3):
    #         self.q_d_lines[i].set_ydata(self.q_d_mat[:, i])
    #         self.q_lines[i].set_ydata(self.q_mat[:, i])
    #     plt.draw()
    #     plt.pause(0.001)
    def cal_foot_place(self):
        global n
        if n != 0:
            pxy0 = [[self.px0], [self.py0]]
            sxy = [[sx[n]], [-(-1) ** n * sy[n]]]
            pxy0 = pxy0 + rotate(n) * sxy

            self.px0 = pxy0[0][0]
            self.py0 = pxy0[1][0]

        """Calculate the coordinate (xbar, ybar)"""
        xybar = [[sx[n + 1] / 2], [-(-1) ** n * sy[n] / 2]]
        xybar = rotate(n + 1) * xybar
        xbar = xybar[0][0]
        ybar = xybar[1][0]
        vxybar = [[(C + 1) / (Tc * S) * xbar], [(C - 1) / (Tc * S) * ybar]]
        vxybar = rotate(0) * vxybar
        vxbar = vxybar[0][0]
        vybar = vxybar[1][0]

        """Target state of CoM of hte n-th step"""
        xd = self.px0 + xbar
        yd = self.py0 + ybar
        vxd = vxbar
        vyd = vybar
        # print(xd)
        # print(vxd)

        """
        Update px, py to be the real foot place in step n,
        a, b are parameters for cost function
        """
        self.px = -a * (C - 1) / D * (xd - C * self.xi - Tc * S * self.vxi) - b * S / (Tc * D) * (
                vxd - S / Tc * self.xi - C * self.vxi)
        self.py = -a * (C - 1) / D * (yd - C * self.yi - Tc * S * self.vyi) - b * S / (Tc * D) * (
                vyd - S / Tc * self.yi - C * self.vyi)
        # print(self.px)

        """Plot out all the foot placements"""
        # plt.plot(px0, py0, 'x')
        # plt.plot(px, py, 'o')

    def cal_CoM_Trajectory(self, t):
        global n

        if n != int(t // Tsup):
            """
            Refresh the steps.
            """
            self.xi = self.xt
            self.yi = self.yt
            self.vxi = self.vxt
            self.vyi = self.vyt

            """Calculate the actual foot place for step n"""
            n = int(t // Tsup)
            if n < np.size(sx):
                self.cal_foot_place()

        self.xt = (self.xi - self.px) * np.cosh((t % Tsup) / Tc) + Tc * self.vxi * np.sinh((t % Tsup) / Tc) + self.px
        self.yt = (self.yi - self.py) * np.cosh((t % Tsup) / Tc) + Tc * self.vyi * np.sinh((t % Tsup) / Tc) + self.py

        self.vxt = (self.xi - self.px) / Tc * np.sinh((t % Tsup) / Tc) + self.vxi * np.cosh((t % Tsup) / Tc)
        self.vyt = (self.yi - self.py) / Tc * np.sinh((t % Tsup) / Tc) + self.vyi * np.cosh((t % Tsup) / Tc)

        print(self.xi)
        print(self.xt)

        return self.xt, self.vxt


if __name__ == '__main__':
    robot = Biped3D()
    robot.run()
